# -*- coding:utf-8 -*-

__author__ = "zxzhao"
# data = 20180123

import os
import shutil
import cv2
import random
import time
import progressbar
import datetime
import  sys
reload(sys)
sys.setdefaultencoding("utf-8")

class batch_processing():
	def __init__(self):
		print 'batch_processing Init()!'

	def get_datetime_now(self):
		cur=datetime.datetime.now()
		return str(cur.month).zfill(2)+str(cur.day).zfill(2)+str(cur.hour).zfill(2)+str(cur.minute).zfill(2)

	def get_file_list(self, file_path, file_type = '', print_path = False):
		file_path = os.path.abspath(file_path)
		path_list = []
		file_list = []
		for rootpath,dirnames,filenames in os.walk(file_path):
			for c_file in filenames:
				if file_type != '' and (not c_file.endswith(file_type)):
					continue
				path_list.append(rootpath)
				file_list.append(c_file)
		if print_path: print (file_path + " : " +str(len(file_list)))
		return path_list, file_list

	def batch_rename(self, path_list, file_list, dataset_name, add_time = False, index_len = 6, file_type = '.jpg', 
					 print_process = True, error_pass = False):
		
		cur_time = self.get_datetime_now()
		file_list_num = len(file_list)

		if not '.' in file_type:
			file_type = '.' + file_type

		if print_process:
			p_progressbar = progressbar.ProgressBar(maxval = file_list_num)
			p_progressbar.start()

		for cur_index in xrange(file_list_num):
			if print_process:
				p_progressbar.update(cur_index+1)

			c_path = path_list[cur_index]
			c_file = file_list[cur_index]
			if add_time:
				save_file = dataset_name + '_' + cur_time + '_' + str(cur_index).zfill(index_len) + file_type
			else:
				save_file = dataset_name + '_' + str(cur_index).zfill(index_len) + file_type
			src_file = os.path.join(c_path, c_file)
			dst_file = os.path.join(c_path, save_file)

			try:
				shutil.move(src_file, dst_file)
			except:
				print ('err_image : ' + src_file)
				if not error_pass:
					return
		if print_process:
			p_progressbar.finish() 

	def batch_resize_img(self, path_list, file_list, file_type = '.jpg', resize_type = 0, 
						 resize_ratio = 1.0, controled_size = 0, print_process = True, error_pass = False):
		# resize_type = 0 : scale up/down the original image at ratio (resize_ratio)
		# resize_type = 1 : scale up/down the original image at ratio, making the maximum edge equal to the specified size (controled_size)
		# resize_type = 2 : scale up/down the original image at ratio, making the minimum edge equal to the specified size (controled_size)

		file_list_num = len(file_list)
		if print_process:
			p_progressbar = progressbar.ProgressBar(maxval = file_list_num)
			p_progressbar.start()
		
		for cur_index in xrange(file_list_num):
			if print_process:
				p_progressbar.update(cur_index+1)

			c_path = path_list[cur_index]
			c_file = file_list[cur_index]
			src_file = os.path.join(c_path, c_file)
			try:
				image_data = cv2.imread(src_file)
				dst_img = self.resize_image_data(image_data, resize_type, resize_ratio, controled_size)

				if not cv2.imwrite(src_file, dst_img):
					print ('no_permission : ' + src_file)
			except:
				print ('err_image : ' + src_file)
				if not error_pass:
					return 
		if print_process:
			p_progressbar.finish() 


	def resize_image_data(self, image_data, resize_type = 0, resize_ratio = 1.0, controled_size = 0):
		# resize_type = 0 : scale up/down the original image at ratio (resize_ratio)
		# resize_type = 1 : scale up/down the original image at ratio, making the maximum edge equal to the specified size (controled_size)
		# resize_type = 2 : scale up/down the original image at ratio, making the minimum edge equal to the specified size (controled_size)
		try:
			if not resize_type in range(0,3) :
				print ("resize_type(between 0 and 3)!")
				return
			if not (resize_ratio > 0 and resize_ratio <= 1):
				print ("resize_ratio(between 0 and 1)!")
				return
			if not (controled_size >= 0):
				print ("controled_size(> 0)!")
				return
		except:
			return
		
		if resize_type == 0 and resize_ratio == 1:
			print ("You need to set resize_ratio(between 0 and 1)!")
			return
		elif resize_type !=0 and controled_size == 0:
			print ("You need to set the specified length(controled_size) greater than 1!")
			return

		img_shape = image_data.shape
		img_width = img_shape[0]
		img_height = img_shape[1]

		if resize_type == 1:
			max_size = max(img_width, img_height)
			resize_ratio = float(controled_size)/float(max_size)
		elif resize_type == 2:
			min_size = min(img_width, img_height)
			resize_ratio = float(controled_size)/float(min_size)
		if resize_ratio > 1.0 :
			resize_ratio = 1.0
		resized_img = cv2.resize(image_data, (int(img_height*resize_ratio), int(img_width*resize_ratio)))
		return resized_img

	def shuffle_list(self, txt_file_path):
		try:
			ouput_file = open(txt_file_path)
			lines = ouput_file.readlines()
			np.random.shuffle(lines)
			if os.path.exists(txt_file_path):
				os.remove(txt_file_path)
			
			with open(txt_file_path,"a+") as output:
				for c_line in lines:
				    output.write(c_line)
				output.close()
		except:
			print ("error :" + txt_file_path)
		return 

	# Get the list of the next level of directory in the current directory
	def get_first_level_dirs(self, base_path):
		folder_list =[]
		cur_list=os.listdir(base_path)
		for c_folder in cur_list:
			if os.path.isdir(os.path.join(base_path, c_folder)):
				folder_list.append(c_folder)
		return folder_list



	# get a list of dataset (train.txt and test.txt)
	def get_label_txt(self, data_basePath, class_label= [], trainval_ratio = 0.8, 
					  file_type = '.jpg', save_relativePath = True, shuffle = True, print_path = False):

		if class_label == []:
			print ("class_label == []")
			return 
		train_list = []
		test_list  = []
		data_basePath = os.path.abspath(data_basePath)
		folder_list = self.get_first_level_dirs(data_basePath)
		for c_folder in folder_list:
			# do every class
			if c_folder in class_label:
				c_class_index = class_label.index(c_folder)
				c_path_list, c_file_list = self.get_file_list(os.path.join(data_basePath, c_folder), 
											file_type = file_type, print_path = print_path)
				
				c_list_num = len(c_file_list)
				train_num  = int(c_list_num * trainval_ratio)

				for cur_index in xrange(c_list_num):
					c_path = c_path_list[cur_index]
					c_file = c_file_list[cur_index]
					src_file = os.path.join(c_path, c_file)
					
					if save_relativePath:
						c_data_basePath = os.path.abspath(os.path.join(data_basePath, "..")) + '/'
						src_file = src_file.replace(c_data_basePath, '', 1)
					cur_line = src_file + ' ' + str(c_class_index) + '\n'
					
					if cur_index < train_num:
						train_list.append(cur_line)
					else:
						test_list.append(cur_line)
					
					if shuffle:
						random.shuffle(train_list)
						random.shuffle(test_list)

		return train_list, test_list

	def write_list_to_file(self, c_list, file_path, rm_existFile = True):

		try:
			if rm_existFile and os.path.exists(file_path):
				os.remove(file_path)
			with open(file_path,"a+") as output:
				for c_line in c_list:
				    output.write(c_line)
				output.close()
		except:
			print ("error :" + txt_file_path)
		return 













